package middleware

import (
	"github.com/gin-gonic/gin"
	"net/http"
	"shortURL/pool"
	"shortURL/util"
	"time"
)

const RATE_PREFIX = "rate:"

// Limiter 定义属性
type Limiter struct {
}

func (l *Limiter) Allow(key string, events int64, per time.Duration) bool {
	conn, err := pool.GetCacheConn("default")
	if err != nil {
		return true
	}
	defer pool.PutCacheConn("default", conn)

	key = RATE_PREFIX + key
	curr := conn.LLen(key).Val()
	if curr >= events {
		return false
	}

	if v := conn.Exists(key).Val(); v == 0 {
		pipe := conn.TxPipeline()
		pipe.RPush(key, key)
		//设置过期时间
		pipe.Expire(key, per)
		_, _ = pipe.Exec()
	} else {
		conn.RPushX(key, key)
	}

	return true
}

func RateMiddleware(c *gin.Context) {
	limiter := &Limiter{}

	events, ok := c.Get("events")
	if !ok {
		events = int64(3)
	}

	per, ok := c.Get("per")
	if !ok {
		per = 1 * time.Minute
	}

	if !limiter.Allow(util.GetIP(), events.(int64), per.(time.Duration)) {
		c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
			"code":    "403",
			"message": "too many request",
		})
		return
	}

	c.Next()
}
